"""
超时控制 (Timeout)

实现超时管理，防止长时间阻塞。

作者: RedFire Team
创建日期: 2025-10-06
版本: v1.0
"""

from typing import Optional, Callable, Any
import asyncio
import logging
from contextlib import asynccontextmanager


class TimeoutError(Exception):
    """超时异常"""
    
    def __init__(self, message: str, timeout: float):
        super().__init__(message)
        self.timeout = timeout


class TimeoutManager:
    """
    超时管理器
    
    提供简单易用的超时控制。
    
    使用示例：
    ```python
    # 方式1：使用execute方法
    timeout_mgr = TimeoutManager(timeout=10.0)
    result = await timeout_mgr.execute(some_async_function, arg1, arg2)
    
    # 方式2：使用上下文管理器
    async with timeout_mgr.context():
        result = await some_async_function()
    ```
    """
    
    def __init__(
        self,
        timeout: float,
        logger: Optional[logging.Logger] = None
    ):
        """
        初始化超时管理器
        
        Args:
            timeout: 超时时间（秒）
            logger: 日志记录器
        """
        self.timeout = timeout
        self.logger = logger or logging.getLogger(__name__)
    
    async def execute(self, func: Callable, *args, **kwargs) -> Any:
        """
        执行带超时控制的函数调用
        
        Args:
            func: 要执行的异步函数
            *args: 位置参数
            **kwargs: 关键字参数
        
        Returns:
            Any: 函数返回值
        
        Raises:
            TimeoutError: 超时时抛出
        """
        try:
            result = await asyncio.wait_for(
                func(*args, **kwargs),
                timeout=self.timeout
            )
            return result
            
        except asyncio.TimeoutError:
            error_msg = f"操作超时: {self.timeout}秒"
            self.logger.error(error_msg)
            raise TimeoutError(error_msg, self.timeout)
    
    @asynccontextmanager
    async def context(self):
        """
        超时上下文管理器
        
        Yields:
            None
        
        Raises:
            TimeoutError: 超时时抛出
        """
        try:
            async with asyncio.timeout(self.timeout):
                yield
        except asyncio.TimeoutError:
            error_msg = f"操作超时: {self.timeout}秒"
            self.logger.error(error_msg)
            raise TimeoutError(error_msg, self.timeout)

